package com.pudding.tools;

import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.output.Response;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.time.LocalDateTime;
import java.util.Arrays;
import java.util.Collections;

public class ToolsDemo2 {

    @Tool("获取当前⽇期")  // 声明⼯具
    public static String dateUtil() {
        return LocalDateTime.now().toString();
    }

    public static void main(String[] args) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
        ChatLanguageModel model = OpenAiChatModel.builder()
                .apiKey(ModelUtil.API_KEY_OPENAI)
                .build();

        //通过⽅法构建⼯具集
        ToolSpecification toolSpecification = ToolSpecifications.toolSpecificationFrom(ToolsDemo2.class.getMethod("dateUtil"));

        UserMessage userMessage = UserMessage.from("今天是⼏⽉⼏号？");
        Response<AiMessage> response = model.generate(Collections.singletonList(userMessage), toolSpecification);
        AiMessage aiMessage = response.content();
        System.out.println(aiMessage);

        // 获取⼯具执⾏结果
        if (aiMessage.hasToolExecutionRequests()) {
            for (ToolExecutionRequest toolExecutionRequest : aiMessage.toolExecutionRequests()) {
                String methodName = toolExecutionRequest.name();
                Method method = ToolsDemo2.class.getMethod(methodName);
                // result就是当前时间
                String result = (String) method.invoke(null);
                System.out.println(result);

                // 构建⼯具执⾏结果消息传递给模型
                ToolExecutionResultMessage toolExecutionResultMessage = ToolExecutionResultMessage.from(toolExecutionRequest.id(), toolExecutionRequest.name(), result);
                AiMessage message = model.generate(Arrays.asList(userMessage, aiMessage, toolExecutionResultMessage)).content();
                System.out.println(message.text());
            }
        }
    }
}
